{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chatglm_cpp\n",
    "from PIL import Image\n",
    "from tool_register import get_tools, run_function, get_tool_desc\n",
    "from loguru import logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline = chatglm_cpp.Pipeline(\"../chatglm3-ggml.bin\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "messages = [chatglm_cpp.ChatMessage(role='system', content=get_tool_desc()) \n",
    "            ,chatglm_cpp.ChatMessage(role='system', content=\"假设你是一个数据库工程师，根据用户的需求提供正确的sql语句\")\n",
    "            , chatglm_cpp.ChatMessage(role='user', content=\"创建一张城市表要求Schema信息如下: Cities(id,city_id,city_name, country, gmt_create, gmt_modify)\")]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "chatMessage = pipeline.chat(\n",
    "                messages,\n",
    "                max_length=2048,\n",
    "                max_context_length=2048,\n",
    "                do_sample=0.95,\n",
    "                top_k=0,\n",
    "                top_p=0.8,\n",
    "                temperature=0.8,\n",
    "                repetition_penalty=1.0,\n",
    "                num_threads=0,\n",
    "                stream=False,\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ChatMessage(role=\"assistant\", content=\"以下是创建城市表的SQL语句： execute_sql\n",
       " ```python\n",
       "tool_call(sql=\\\"CREATE TABLE Cities (id INTEGER PRIMARY KEY, city_id INTEGER, city_name VARCHAR(50), country VARCHAR(50), gmt_create TIMESTAMP, gmt_modify TIMESTAMP);\\\")\n",
       "```\", tool_calls=[ToolCallMessage(type=\"function\", function=FunctionMessage(name=\"\", arguments=\"tool_call(sql=\\\"CREATE TABLE Cities (id INTEGER PRIMARY KEY, city_id INTEGER, city_name VARCHAR(50), country VARCHAR(50), gmt_create TIMESTAMP, gmt_modify TIMESTAMP);\\\")\"), code=CodeMessage(input=\"\"))])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chatMessage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages.append(chatglm_cpp.ChatMessage(role=chatMessage.role, content=chatMessage.content))\n",
    "messages.append(chatglm_cpp.ChatMessage(role=\"user\", content=\"查询一下城市在北京的用户数，你需要首先生成一个sql语句，然后调用执行sql工具获取结果\"))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "chatMessage = pipeline.chat(\n",
    "                messages,\n",
    "                max_length=2048,\n",
    "                max_context_length=2048,\n",
    "                do_sample=0.95,\n",
    "                top_k=0,\n",
    "                top_p=0.8,\n",
    "                temperature=0.8,\n",
    "                repetition_penalty=1.0,\n",
    "                num_threads=0,\n",
    "                stream=False,\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ChatMessage(role=\"assistant\", content=\"以下是查询城市在北京的用户数的SQL语句： execute_sql\n",
       " ```python\n",
       "tool_call(sql=\\\"SELECT COUNT(*) FROM Cities WHERE city_name = '北京';\\\")\n",
       "```\", tool_calls=[ToolCallMessage(type=\"function\", function=FunctionMessage(name=\"\", arguments=\"tool_call(sql=\\\"SELECT COUNT(*) FROM Cities WHERE city_name = '北京';\\\")\"), code=CodeMessage(input=\"\"))])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chatMessage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "chatglm3-demo",
   "language": "python",
   "name": "chatglm3-demo"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
